#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
pip install wheel
pip install scapy
"""
from scapy.all import *
from collections import defaultdict
import json
from datetime import datetime
import sys
import os

class RedisPcapParser:
    def __init__(self, pcap_file, redis_port=6379):
        self.pcap_file = pcap_file
        self.redis_port = redis_port
        self.commands_by_client = defaultdict(list)
        self.current_streams = {}

    def parse(self):
        """解析pcap文件"""
        packets = rdpcap(self.pcap_file)
        
        for pkt in packets:
            if TCP in pkt and pkt[TCP].dport == self.redis_port and IP in pkt:
                self._process_packet(pkt)
        
        # 为每个客户端的命令按时间排序
        for client_ip in self.commands_by_client:
            self.commands_by_client[client_ip].sort(key=lambda x: x['timestamp'])

    def _process_packet(self, pkt):
        """处理单个数据包"""
        stream_key = (pkt[IP].src, pkt[IP].dst, pkt[TCP].sport, pkt[TCP].dport)
        
        if stream_key not in self.current_streams:
            self.current_streams[stream_key] = {
                'client_buffer': b'',
                'client_ip': pkt[IP].src,
                'server_ip': pkt[IP].dst
            }
        
        stream = self.current_streams[stream_key]
        payload = bytes(pkt[TCP].payload)
        timestamp = datetime.fromtimestamp(pkt.time)
        
        if pkt[TCP].sport != self.redis_port:
            stream['client_buffer'] += payload
            self._parse_redis_protocol(
                stream['client_buffer'], 
                timestamp,
                stream['client_ip'],
                stream['server_ip']
            )
            
            if len(stream['client_buffer']) > 1024:
                stream['client_buffer'] = b''

    def _parse_redis_protocol(self, data, timestamp, client_ip, server_ip):
        """解析Redis协议"""
        try:
            if not data.startswith(b'*'):
                return
                
            parts = data.split(b'\r\n')
            if len(parts) < 3:
                return
                
            cmd_parts = []
            i = 0
            while i < len(parts):
                if parts[i].startswith(b'*'):
                    try:
                        array_len = int(parts[i][1:])
                    except ValueError:
                        break
                    i += 1
                    for _ in range(array_len):
                        if i < len(parts) and parts[i].startswith(b'$'):
                            try:
                                param_len = int(parts[i][1:])
                            except ValueError:
                                break
                            i += 1
                            if i < len(parts):
                                try:
                                    param = parts[i][:param_len].decode('utf-8', errors='replace')
                                except UnicodeDecodeError:
                                    param = str(parts[i][:param_len])
                                cmd_parts.append(param)
                                i += 1
                else:
                    i += 1
            
            if not cmd_parts:
                return
                
            cmd = cmd_parts[0].upper()
            args = cmd_parts[1:]
            
            self.commands_by_client[client_ip].append({
                'timestamp': timestamp.isoformat(),
                'command': cmd,
                'args': args,
                'server_ip': server_ip,
                'raw': data.decode('ascii', errors='replace')[:200]
            })
            
        except Exception as e:
            print("Error parsing Redis command: {}".format(e))

    def save_results(self, output_file):
        """保存结果到JSON文件"""
        result = {
            'clients': {},
            'summary': {
                'total_clients': len(self.commands_by_client),
                'total_commands': sum(len(cmds) for cmds in self.commands_by_client.values()),
                'start_time': min(cmd['timestamp'] for cmds in self.commands_by_client.values() for cmd in cmds) 
                              if self.commands_by_client else None,
                'end_time': max(cmd['timestamp'] for cmds in self.commands_by_client.values() for cmd in cmds) 
                            if self.commands_by_client else None,
                'redis_port': self.redis_port
            }
        }
        
        for client_ip, commands in self.commands_by_client.items():
            result['clients'][client_ip] = {
                'total_commands': len(commands),
                'first_command': commands[0]['timestamp'] if commands else None,
                'last_command': commands[-1]['timestamp'] if commands else None,
                'commands': commands
            }
        
        with open(output_file, 'w') as f:
            json.dump(result, f, indent=2, ensure_ascii=False)
        
        print("Results saved to {}".format(output_file))
        print("Total clients: {}".format(result['summary']['total_clients']))
        print("Total commands: {}".format(result['summary']['total_commands']))

def main():
    if len(sys.argv) < 2:
        print("Usage: {} <pcap_file> [output_json] [redis_port]".format(sys.argv[0]))
        sys.exit(1)
    
    pcap_file = sys.argv[1]
    if not os.path.exists(pcap_file):
        print("Error: PCAP file not found: {}".format(pcap_file))
        sys.exit(1)
    
    output_file = sys.argv[2] if len(sys.argv) > 2 else "redis_commands.json"
    redis_port = int(sys.argv[3]) if len(sys.argv) > 3 else 6379
    
    print("Parsing Redis traffic from {} (port: {})...".format(pcap_file, redis_port))
    parser = RedisPcapParser(pcap_file, redis_port)
    parser.parse()
    parser.save_results(output_file)
    
    # 打印示例数据
    print("\nSample data:")
    for client_ip, data in parser.commands_by_client.items():
        print("\nClient {} ({} commands):".format(client_ip, len(data)))
        for cmd in data[:3]:
            print("  [{}] {} {}".format(cmd['timestamp'], cmd['command'], ' '.join(cmd['args'])))
        if len(data) > 3:
            print("  ... and {} more commands".format(len(data)-3))

if __name__ == "__main__":
    main()
